class DigitalHumanGenerator:
    """数字人视频生成工具类"""

    def __init__(self, config: DigitalHumanConfig):
        self.config = config
        self.supported_providers = ["sadtalker", "heygen", "synthesia", "d_id"]

    def create_talking_head(self,
                            text: str,
                            audio_path: Optional[str] = None,
                            image_path: Optional[str] = None,
                            avatar_id: str = "default",
                            provider: str = "sadtalker",
                            **kwargs) -> Dict[str, Any]:
        """生成说话头像视频"""

        if provider == "sadtalker":
            return self._sadtalker_generate(text, audio_path, image_path, **kwargs)
        elif provider == "heygen":
            return self._heygen_generate(text, avatar_id, **kwargs)
        elif provider == "synthesia":
            return self._synthesia_generate(text, avatar_id, **kwargs)
        elif provider == "d_id":
            return self._d_id_generate(text, image_path, **kwargs)
        else:
            raise ValueError(f"不支持的提供商: {provider}")

    def _sadtalker_generate(self, text: str, audio_path: Optional[str], image_path: Optional[str], **kwargs) -> Dict:
        """使用SadTalker生成"""
        api_key = self.config.configs["sadtalker"]["api_key"]
        base_url = self.config.configs["sadtalker"]["base_url"]

        # 准备请求数据
        data = {
            "text": text,
            "preprocess": "full",  # 完整预处理
            "still_mode": True,  # 静态模式
            "enhancer": True,  # 启用增强
            "expression_scale": 1.0,
            **kwargs
        }

        files = {}
        if audio_path and os.path.exists(audio_path):
            files['audio'] = open(audio_path, 'rb')
        if image_path and os.path.exists(image_path):
            files['image'] = open(image_path, 'rb')

        try:
            response = requests.post(
                f"{base_url}/generate",
                headers={"Authorization": f"Bearer {api_key}"},
                data=data,
                files=files if files else None,
                timeout=300  # 5分钟超时
            )

            if response.status_code == 200:
                result = response.json()
                return {
                    "success": True,
                    "video_url": result.get("video_url"),
                    "duration": result.get("duration"),
                    "provider": "sadtalker"
                }
            else:
                return {
                    "success": False,
                    "error": f"生成失败: {response.text}",
                    "provider": "sadtalker"
                }

        except Exception as e:
            return {
                "success": False,
                "error": f"请求异常: {str(e)}",
                "provider": "sadtalker"
            }
        finally:
            # 关闭文件
            for file in files.values():
                file.close()

    def _heygen_generate(self, text: str, avatar_id: str, **kwargs) -> Dict:
        """使用HeyGen生成"""
        api_key = self.config.configs["heygen"]["api_key"]
        base_url = self.config.configs["heygen"]["base_url"]

        payload = {
            "video_inputs": [{
                "character": {
                    "type": "avatar",
                    "avatar_id": avatar_id,
                    "avatar_style": "normal"
                },
                "voice": {
                    "type": "text",
                    "input_text": text,
                    "voice_id": "1bd001e7e50f421d891986aad5158bc8"  # 中文声音
                },
                "background": {
                    "type": "color",
                    "value": "#FFFFFF"
                }
            }],
            "dimension": {"width": 1920, "height": 1080},
            **kwargs
        }

        try:
            # 创建任务
            response = requests.post(
                f"{base_url}/video/generate",
                headers={
                    "X-Api-Key": api_key,
                    "Content-Type": "application/json"
                },
                json=payload,
                timeout=60
            )

            if response.status_code == 201:
                task_id = response.json().get("data", {}).get("video_id")
                return {
                    "success": True,
                    "task_id": task_id,
                    "provider": "heygen",
                    "status_url": f"{base_url}/video/{task_id}"
                }
            else:
                return {
                    "success": False,
                    "error": f"任务创建失败: {response.text}",
                    "provider": "heygen"
                }

        except Exception as e:
            return {
                "success": False,
                "error": f"请求异常: {str(e)}",
                "provider": "heygen"
            }

    def _synthesia_generate(self, text: str, avatar_id: str, **kwargs) -> Dict:
        """使用Synthesia生成"""
        api_key = self.config.configs["synthesia"]["api_key"]
        base_url = self.config.configs["synthesia"]["base_url"]

        payload = {
            "test": False,
            "input": [
                {
                    "avatar": avatar_id,
                    "text": text,
                    "script": {
                        "type": "text",
                        "input": text
                    },
                    "background": {
                        "type": "color",
                        "value": "#FFFFFF"
                    }
                }
            ],
            **kwargs
        }

        try:
            response = requests.post(
                f"{base_url}/videos",
                headers={"Authorization": api_key},
                json=payload,
                timeout=60
            )

            if response.status_code == 201:
                video_id = response.json().get("id")
                return {
                    "success": True,
                    "video_id": video_id,
                    "provider": "synthesia",
                    "status": "processing"
                }
            else:
                return {
                    "success": False,
                    "error": f"生成失败: {response.text}",
                    "provider": "synthesia"
                }

        except Exception as e:
            return {
                "success": False,
                "error": f"请求异常: {str(e)}",
                "provider": "synthesia"
            }

    def _d_id_generate(self, text: str, image_path: Optional[str], **kwargs) -> Dict:
        """使用D-ID生成"""
        api_key = self.config.configs["d_id"]["api_key"]
        base_url = self.config.configs["d_id"]["base_url"]

        # 如果有图片，先上传
        source_url = None
        if image_path and os.path.exists(image_path):
            with open(image_path, "rb") as f:
                upload_response = requests.post(
                    f"{base_url}/images",
                    headers={"Authorization": f"Basic {api_key}"},
                    files={"image": f}
                )
                if upload_response.status_code == 200:
                    source_url = upload_response.json().get("url")

        if not source_url:
            # 使用默认图片或返回错误
            return {
                "success": False,
                "error": "需要提供源图片",
                "provider": "d_id"
            }

        # 创建说话视频
        payload = {
            "source_url": source_url,
            "script": {
                "type": "text",
                "input": text,
                "provider": {
                    "type": "microsoft",
                    "voice_id": "zh-CN-XiaoxiaoNeural"
                }
            },
            "config": {
                "fluent": "true",
                "pad_audio": "0.0"
            },
            **kwargs
        }

        try:
            response = requests.post(
                f"{base_url}/talks",
                headers={
                    "Authorization": f"Basic {api_key}",
                    "Content-Type": "application/json"
                },
                json=payload,
                timeout=60
            )

            if response.status_code == 201:
                talk_id = response.json().get("id")
                return {
                    "success": True,
                    "talk_id": talk_id,
                    "provider": "d_id",
                    "status_url": f"{base_url}/talks/{talk_id}"
                }
            else:
                return {
                    "success": False,
                    "error": f"生成失败: {response.text}",
                    "provider": "d_id"
                }

        except Exception as e:
            return {
                "success": False,
                "error": f"请求异常: {str(e)}",
                "provider": "d_id"
            }

    def check_video_status(self, provider: str, task_id: str) -> Dict:
        """检查视频生成状态"""
        if provider == "heygen":
            return self._check_heygen_status(task_id)
        elif provider == "synthesia":
            return self._check_synthesia_status(task_id)
        elif provider == "d_id":
            return self._check_d_id_status(task_id)
        else:
            return {"success": False, "error": f"不支持的提供商: {provider}"}

    def _check_heygen_status(self, task_id: str) -> Dict:
        """检查HeyGen任务状态"""
        api_key = self.config.configs["heygen"]["api_key"]
        base_url = self.config.configs["heygen"]["base_url"]

        try:
            response = requests.get(
                f"{base_url}/video/{task_id}",
                headers={"X-Api-Key": api_key}
            )

            if response.status_code == 200:
                data = response.json().get("data", {})
                status = data.get("status")
                result = {
                    "success": True,
                    "status": status,
                    "provider": "heygen"
                }

                if status == "completed":
                    result["video_url"] = data.get("video_url")
                    result["duration"] = data.get("duration")

                return result
            else:
                return {
                    "success": False,
                    "error": f"状态查询失败: {response.text}",
                    "provider": "heygen"
                }

        except Exception as e:
            return {
                "success": False,
                "error": f"查询异常: {str(e)}",
                "provider": "heygen"
            }